{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import os\n",
    "os.sys.path.append('..')\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1\"\n",
    "import torch\n",
    "from torch.autograd import Variable\n",
    "from torchvision import datasets, transforms\n",
    "from scipy.misc import imsave\n",
    "import scipy.io\n",
    "import warnings\n",
    "import sys\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.misc\n",
    "\n",
    "from darknet_multi import Darknet\n",
    "from utils import *\n",
    "import dataset_multi\n",
    "from MeshPly import MeshPly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2018-05-06 14:09:50 Testing ape...\n",
      "2018-05-06 14:10:15    Acc using 5 px 2D Projection = 7.01%\n",
      "2018-05-06 14:10:15    Acc using 10 px 2D Projection = 40.43%\n",
      "2018-05-06 14:10:15    Acc using 15 px 2D Projection = 59.83%\n",
      "2018-05-06 14:10:15    Acc using 20 px 2D Projection = 68.55%\n",
      "2018-05-06 14:10:15    Acc using 25 px 2D Projection = 72.05%\n",
      "2018-05-06 14:10:15    Acc using 30 px 2D Projection = 73.68%\n",
      "2018-05-06 14:10:15    Acc using 35 px 2D Projection = 74.53%\n",
      "2018-05-06 14:10:15    Acc using 40 px 2D Projection = 75.13%\n",
      "2018-05-06 14:10:15    Acc using 45 px 2D Projection = 75.73%\n",
      "2018-05-06 14:10:15    Acc using 50 px 2D Projection = 76.50%\n",
      "2018-05-06 14:10:18 Testing can...\n",
      "2018-05-06 14:10:47    Acc using 5 px 2D Projection = 11.18%\n",
      "2018-05-06 14:10:47    Acc using 10 px 2D Projection = 57.83%\n",
      "2018-05-06 14:10:47    Acc using 15 px 2D Projection = 79.95%\n",
      "2018-05-06 14:10:47    Acc using 20 px 2D Projection = 85.75%\n",
      "2018-05-06 14:10:47    Acc using 25 px 2D Projection = 88.73%\n",
      "2018-05-06 14:10:47    Acc using 30 px 2D Projection = 90.39%\n",
      "2018-05-06 14:10:47    Acc using 35 px 2D Projection = 91.80%\n",
      "2018-05-06 14:10:47    Acc using 40 px 2D Projection = 93.21%\n",
      "2018-05-06 14:10:47    Acc using 45 px 2D Projection = 93.62%\n",
      "2018-05-06 14:10:47    Acc using 50 px 2D Projection = 93.79%\n",
      "2018-05-06 14:10:50 Testing cat...\n",
      "2018-05-06 14:11:16    Acc using 5 px 2D Projection = 3.62%\n",
      "2018-05-06 14:11:16    Acc using 10 px 2D Projection = 23.25%\n",
      "2018-05-06 14:11:16    Acc using 15 px 2D Projection = 39.51%\n",
      "2018-05-06 14:11:16    Acc using 20 px 2D Projection = 49.45%\n",
      "2018-05-06 14:11:16    Acc using 25 px 2D Projection = 54.76%\n",
      "2018-05-06 14:11:16    Acc using 30 px 2D Projection = 57.96%\n",
      "2018-05-06 14:11:16    Acc using 35 px 2D Projection = 59.56%\n",
      "2018-05-06 14:11:16    Acc using 40 px 2D Projection = 60.99%\n",
      "2018-05-06 14:11:16    Acc using 45 px 2D Projection = 62.51%\n",
      "2018-05-06 14:11:16    Acc using 50 px 2D Projection = 63.27%\n",
      "2018-05-06 14:11:19 Testing duck...\n",
      "2018-05-06 14:11:42    Acc using 5 px 2D Projection = 5.07%\n",
      "2018-05-06 14:11:42    Acc using 10 px 2D Projection = 18.20%\n",
      "2018-05-06 14:11:42    Acc using 15 px 2D Projection = 30.88%\n",
      "2018-05-06 14:11:42    Acc using 20 px 2D Projection = 55.12%\n",
      "2018-05-06 14:11:42    Acc using 25 px 2D Projection = 75.15%\n",
      "2018-05-06 14:11:42    Acc using 30 px 2D Projection = 81.45%\n",
      "2018-05-06 14:11:42    Acc using 35 px 2D Projection = 83.20%\n",
      "2018-05-06 14:11:42    Acc using 40 px 2D Projection = 83.64%\n",
      "2018-05-06 14:11:42    Acc using 45 px 2D Projection = 83.90%\n",
      "2018-05-06 14:11:42    Acc using 50 px 2D Projection = 84.16%\n",
      "2018-05-06 14:11:45 Testing driller...\n",
      "2018-05-06 14:12:10    Acc using 5 px 2D Projection = 1.40%\n",
      "2018-05-06 14:12:10    Acc using 10 px 2D Projection = 17.38%\n",
      "2018-05-06 14:12:10    Acc using 15 px 2D Projection = 39.87%\n",
      "2018-05-06 14:12:10    Acc using 20 px 2D Projection = 62.93%\n",
      "2018-05-06 14:12:10    Acc using 25 px 2D Projection = 80.64%\n",
      "2018-05-06 14:12:10    Acc using 30 px 2D Projection = 89.87%\n",
      "2018-05-06 14:12:10    Acc using 35 px 2D Projection = 94.89%\n",
      "2018-05-06 14:12:10    Acc using 40 px 2D Projection = 95.88%\n",
      "2018-05-06 14:12:10    Acc using 45 px 2D Projection = 96.54%\n",
      "2018-05-06 14:12:10    Acc using 50 px 2D Projection = 96.87%\n",
      "2018-05-06 14:12:13 Testing glue...\n",
      "2018-05-06 14:12:31    Acc using 5 px 2D Projection = 6.53%\n",
      "2018-05-06 14:12:31    Acc using 10 px 2D Projection = 26.91%\n",
      "2018-05-06 14:12:31    Acc using 15 px 2D Projection = 39.65%\n",
      "2018-05-06 14:12:31    Acc using 20 px 2D Projection = 46.18%\n",
      "2018-05-06 14:12:31    Acc using 25 px 2D Projection = 49.50%\n",
      "2018-05-06 14:12:31    Acc using 30 px 2D Projection = 51.83%\n",
      "2018-05-06 14:12:31    Acc using 35 px 2D Projection = 53.05%\n",
      "2018-05-06 14:12:31    Acc using 40 px 2D Projection = 53.16%\n",
      "2018-05-06 14:12:31    Acc using 45 px 2D Projection = 53.93%\n",
      "2018-05-06 14:12:31    Acc using 50 px 2D Projection = 54.71%\n",
      "2018-05-06 14:12:45 Testing holepuncher...\n",
      "2018-05-06 14:19:31    Acc using 5 px 2D Projection = 8.26%\n",
      "2018-05-06 14:19:31    Acc using 10 px 2D Projection = 39.50%\n",
      "2018-05-06 14:19:31    Acc using 15 px 2D Projection = 53.31%\n",
      "2018-05-06 14:19:31    Acc using 20 px 2D Projection = 62.56%\n",
      "2018-05-06 14:19:31    Acc using 25 px 2D Projection = 68.02%\n",
      "2018-05-06 14:19:31    Acc using 30 px 2D Projection = 74.71%\n",
      "2018-05-06 14:19:31    Acc using 35 px 2D Projection = 80.74%\n",
      "2018-05-06 14:19:31    Acc using 40 px 2D Projection = 85.62%\n",
      "2018-05-06 14:19:31    Acc using 45 px 2D Projection = 89.59%\n",
      "2018-05-06 14:19:31    Acc using 50 px 2D Projection = 91.49%\n"
     ]
    }
   ],
   "source": [
    "def valid(datacfg, cfgfile, weightfile, conf_th):\n",
    "    def truths_length(truths):\n",
    "        for i in range(50):\n",
    "            if truths[i][1] == 0:\n",
    "                return i\n",
    "\n",
    "    # Parse configuration files\n",
    "    options       = read_data_cfg(datacfg)\n",
    "    valid_images  = options['valid']\n",
    "    meshname      = options['mesh']\n",
    "    backupdir     = options['backup']\n",
    "    name          = options['name']\n",
    "    prefix        = 'results'\n",
    "    # Read object model information, get 3D bounding box corners\n",
    "    mesh          = MeshPly(meshname)\n",
    "    vertices      = np.c_[np.array(mesh.vertices), np.ones((len(mesh.vertices), 1))].transpose()\n",
    "    corners3D     = get_3D_corners(vertices)\n",
    "    # Read intrinsic camera parameters\n",
    "    internal_calibration = get_camera_intrinsic()\n",
    "\n",
    "    # Get validation file names\n",
    "    with open(valid_images) as fp:\n",
    "        tmp_files = fp.readlines()\n",
    "        valid_files = [item.rstrip() for item in tmp_files]\n",
    "    \n",
    "    # Specicy model, load pretrained weights, pass to GPU and set the module in evaluation mode\n",
    "    model = Darknet(cfgfile)\n",
    "    model.load_weights(weightfile)\n",
    "    model.cuda()\n",
    "    model.eval()\n",
    "\n",
    "    # Get the parser for the test dataset\n",
    "    valid_dataset = dataset_multi.listDataset(valid_images, shape=(model.width, model.height),\n",
    "                       shuffle=False,\n",
    "                       objclass=name,\n",
    "                       transform=transforms.Compose([\n",
    "                           transforms.ToTensor(),\n",
    "                       ]))\n",
    "    valid_batchsize = 1\n",
    "\n",
    "    # Specify the number of workers for multiple processing, get the dataloader for the test dataset\n",
    "    kwargs = {'num_workers': 4, 'pin_memory': True}\n",
    "    test_loader = torch.utils.data.DataLoader(\n",
    "        valid_dataset, batch_size=valid_batchsize, shuffle=False, **kwargs) \n",
    "\n",
    "    # Parameters\n",
    "    visualize       = False\n",
    "    use_cuda        = True\n",
    "    num_classes     = 13\n",
    "    anchors         = [1.4820, 2.2412, 2.0501, 3.1265, 2.3946, 4.6891, 3.1018, 3.9910, 3.4879, 5.8851]\n",
    "    num_anchors     = 5\n",
    "    eps             = 1e-5\n",
    "    conf_thresh     = conf_th\n",
    "    iou_thresh      = 0.5\n",
    "\n",
    "    # Parameters to save\n",
    "    errs_2d = []\n",
    "    edges = [[1, 2], [1, 3], [1, 5], [2, 4], [2, 6], [3, 4], [3, 7], [4, 8], [5, 6], [5, 7], [6, 8], [7, 8]]\n",
    "    edges_corners = [[0, 1], [0, 2], [0, 4], [1, 3], [1, 5], [2, 3], [2, 6], [3, 7], [4, 5], [4, 6], [5, 7], [6, 7]]\n",
    "\n",
    "    # Iterate through test batches (Batch size for test data is 1)\n",
    "    count = 0\n",
    "    logging('Testing {}...'.format(name))\n",
    "    for batch_idx, (data, target) in enumerate(test_loader):\n",
    "        \n",
    "        # Images\n",
    "        img = data[0, :, :, :]\n",
    "        img = img.numpy().squeeze()\n",
    "        img = np.transpose(img, (1, 2, 0))\n",
    "        \n",
    "        t1 = time.time()\n",
    "        # Pass data to GPU\n",
    "        if use_cuda:\n",
    "            data = data.cuda()\n",
    "            target = target.cuda()\n",
    "        \n",
    "        # Wrap tensors in Variable class, set volatile=True for inference mode and to use minimal memory during inference\n",
    "        data = Variable(data, volatile=True)\n",
    "        t2 = time.time()\n",
    "        \n",
    "        # Forward pass\n",
    "        output = model(data).data  \n",
    "        t3 = time.time()\n",
    "        \n",
    "        # Using confidence threshold, eliminate low-confidence predictions\n",
    "        trgt = target[0].view(-1, 21)\n",
    "        all_boxes = get_corresponding_region_boxes(output, conf_thresh, num_classes, anchors, num_anchors, int(trgt[0][0]), only_objectness=0)        \n",
    "        t4 = time.time()\n",
    "        \n",
    "        # Iterate through all images in the batch\n",
    "        for i in range(output.size(0)):\n",
    "            \n",
    "            # For each image, get all the predictions\n",
    "            boxes   = all_boxes[i]\n",
    "            \n",
    "            # For each image, get all the targets (for multiple object pose estimation, there might be more than 1 target per image)\n",
    "            truths  = target[i].view(-1, 21)\n",
    "            \n",
    "            # Get how many object are present in the scene\n",
    "            num_gts = truths_length(truths)\n",
    "\n",
    "            # Iterate through each ground-truth object\n",
    "            for k in range(num_gts):\n",
    "                box_gt        = [truths[k][1], truths[k][2], truths[k][3], truths[k][4], truths[k][5], truths[k][6], \n",
    "                                truths[k][7], truths[k][8], truths[k][9], truths[k][10], truths[k][11], truths[k][12], \n",
    "                                truths[k][13], truths[k][14], truths[k][15], truths[k][16], truths[k][17], truths[k][18], 1.0, 1.0, truths[k][0]]\n",
    "                best_conf_est = -1\n",
    "                \n",
    "\n",
    "                # If the prediction has the highest confidence, choose it as our prediction\n",
    "                for j in range(len(boxes)):\n",
    "                    if (boxes[j][18] > best_conf_est) and (boxes[j][20] == int(truths[k][0])):\n",
    "                        best_conf_est = boxes[j][18]\n",
    "                        box_pr        = boxes[j]\n",
    "                        bb2d_gt       = get_2d_bb(box_gt[:18], output.size(3))\n",
    "                        bb2d_pr       = get_2d_bb(box_pr[:18], output.size(3))\n",
    "                        iou           = bbox_iou(bb2d_gt, bb2d_pr)\n",
    "                        match         = corner_confidence9(box_gt[:18], torch.FloatTensor(boxes[j][:18]))\n",
    "                    \n",
    "                # Denormalize the corner predictions \n",
    "                corners2D_gt = np.array(np.reshape(box_gt[:18], [9, 2]), dtype='float32')\n",
    "                corners2D_pr = np.array(np.reshape(box_pr[:18], [9, 2]), dtype='float32')\n",
    "                corners2D_gt[:, 0] = corners2D_gt[:, 0] * 640\n",
    "                corners2D_gt[:, 1] = corners2D_gt[:, 1] * 480               \n",
    "                corners2D_pr[:, 0] = corners2D_pr[:, 0] * 640\n",
    "                corners2D_pr[:, 1] = corners2D_pr[:, 1] * 480\n",
    "                corners2D_gt_corrected = fix_corner_order(corners2D_gt) # Fix the order of the corners in OCCLUSION\n",
    "                \n",
    "                # Compute [R|t] by pnp\n",
    "                objpoints3D = np.array(np.transpose(np.concatenate((np.zeros((3, 1)), corners3D[:3, :]), axis=1)), dtype='float32')\n",
    "                K = np.array(internal_calibration, dtype='float32')\n",
    "                R_gt, t_gt = pnp(objpoints3D,  corners2D_gt_corrected, K)\n",
    "                R_pr, t_pr = pnp(objpoints3D,  corners2D_pr, K)\n",
    "                \n",
    "                # Compute pixel error\n",
    "                Rt_gt        = np.concatenate((R_gt, t_gt), axis=1)\n",
    "                Rt_pr        = np.concatenate((R_pr, t_pr), axis=1)\n",
    "                proj_2d_gt   = compute_projection(vertices, Rt_gt, internal_calibration) \n",
    "                proj_2d_pred = compute_projection(vertices, Rt_pr, internal_calibration) \n",
    "                proj_corners_gt = np.transpose(compute_projection(corners3D, Rt_gt, internal_calibration)) \n",
    "                proj_corners_pr = np.transpose(compute_projection(corners3D, Rt_pr, internal_calibration)) \n",
    "                norm         = np.linalg.norm(proj_2d_gt - proj_2d_pred, axis=0)\n",
    "                pixel_dist   = np.mean(norm)\n",
    "                errs_2d.append(pixel_dist)\n",
    "\n",
    "                \n",
    "                if visualize:\n",
    "                    # Visualize\n",
    "                    plt.xlim((0, 640))\n",
    "                    plt.ylim((0, 480))\n",
    "                    plt.imshow(scipy.misc.imresize(img, (480, 640)))\n",
    "                    # Projections\n",
    "                    for edge in edges_corners:\n",
    "                        plt.plot(proj_corners_gt[edge, 0], proj_corners_gt[edge, 1], color='g', linewidth=3.0)\n",
    "                        plt.plot(proj_corners_pr[edge, 0], proj_corners_pr[edge, 1], color='b', linewidth=3.0)\n",
    "                    plt.gca().invert_yaxis()\n",
    "                    plt.show()\n",
    "\n",
    "        t5 = time.time()\n",
    "\n",
    "    # Compute 2D projection score\n",
    "    for px_threshold in [5, 10, 15, 20, 25, 30, 35, 40, 45, 50]:\n",
    "        acc = len(np.where(np.array(errs_2d) <= px_threshold)[0]) * 100. / (len(errs_2d)+eps)\n",
    "        # Print test statistics\n",
    "        logging('   Acc using {} px 2D Projection = {:.2f}%'.format(px_threshold, acc))\n",
    "\n",
    "conf_th = 0.05\n",
    "cfgfile = 'cfg/yolo-pose-multi.cfg'\n",
    "weightfile = 'backup_multi/model_backup2.weights'\n",
    "datacfg = 'cfg/ape_occlusion.data'\n",
    "valid(datacfg, cfgfile, weightfile, conf_th)\n",
    "datacfg = 'cfg/can_occlusion.data'\n",
    "valid(datacfg, cfgfile, weightfile, conf_th)\n",
    "datacfg = 'cfg/cat_occlusion.data'\n",
    "valid(datacfg, cfgfile, weightfile, conf_th)\n",
    "datacfg = 'cfg/duck_occlusion.data'\n",
    "valid(datacfg, cfgfile, weightfile, conf_th)\n",
    "datacfg = 'cfg/driller_occlusion.data'\n",
    "valid(datacfg, cfgfile, weightfile, conf_th)\n",
    "datacfg = 'cfg/glue_occlusion.data'\n",
    "valid(datacfg, cfgfile, weightfile, conf_th)\n",
    "datacfg = 'cfg/holepuncher_occlusion.data'\n",
    "valid(datacfg, cfgfile, weightfile, conf_th)\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
